package top.v5it.japi.plus.core.spring.web.interceptor;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.servlet.ServletUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.HandlerInterceptor;
import top.v5it.japi.plus.common.context.TraceContextHolder;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Optional;

import static top.v5it.japi.plus.common.ConstantCommon.*;

/**
 * 链路追踪拦截器
 *
 * @author zhanpu
 * @date 2022/8/26
 */
@Slf4j
public class TraceInterceptor implements HandlerInterceptor {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        final String traceId = Optional
                .ofNullable(request.getHeader(HEAD_TRACE_ID_NAME))
                .filter(StrUtil::isNotBlank)
                .orElse(IdUtil.randomUUID());
        String spanId = request.getHeader(HEAD_SPAN_ID_NAME);
        spanId = StrUtil.isBlank(spanId) ? ROOT_SPAN_ID : TraceContextHolder.genSpanId(spanId);
        final String ip = ServletUtil.getClientIP(request);
        TraceContextHolder.putIp(ip);
        TraceContextHolder.putTraceId(traceId);
        TraceContextHolder.putSpanId(spanId);
        TraceContextHolder.start();
        final String name = getTraceName(request);
        log.info("{}链路追踪开始, traceId:{}, spanId:{}, ip:{}", name, traceId, spanId, ip);
        return true;
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        final long costTime = TraceContextHolder.stopAndGet();
        final String ip = TraceContextHolder.getIp();
        final String name = getTraceName(request);
        log.info("{}链路追踪结束, ip:{}, 用时:[{}]毫秒", name, ip, costTime);
        TraceContextHolder.removeIp();
        TraceContextHolder.removeTraceId();
        TraceContextHolder.removeSpanId();
    }

    private String getTraceName(HttpServletRequest request) {
        return Optional
                .ofNullable(System.getProperty(SERVER_NAME))
                .filter(StrUtil::isNotBlank)
                .orElse("接口URI:[" + request.getRequestURI() + "]");
    }
}
